import torch.nn as nn
import torch.nn.functional as F


class ResNet18(nn.Module):
    def __init__(self):
        super(ResNet18, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(True)
        )
        self.block1 = nn.Sequential(
            Basic_block(64),
            Basic_block(64),
            Basic_block(64)
        )
        self.block2 = nn.Sequential(
            Basic_Block1(64, 128),
            Basic_block(128),
            Basic_block(128)
        )
        self.block3 = nn.Sequential(
            Basic_Block1(128, 256),
            Basic_block(256),
            Basic_block(256)
        )
        self.avg_POOL = nn.AvgPool2d(8)
        self.fc = nn.Linear(256, 10)

    def forward(self, x):
        y = self.conv1(x)
        y = self.block1(y)
        y = self.block2(y)
        y = self.block3(y)
        y = self.avg_POOL(y)
        y = y.view(x.size(0), -1)
        y = self.fc(y)
        return y


# 实线部分
class Basic_block(nn.Module):
    def __init__(self, in_channels):
        super(Basic_block, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1, stride=1),
            nn.BatchNorm2d(in_channels),
            nn.ReLU(True)
        )

        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1, stride=1),
            nn.BatchNorm2d(in_channels)
        )

    def forward(self, x):
        y = self.conv1(x)
        y = self.conv2(y)
        y = F.relu(y + x)
        return y


# 虚线部分
class Basic_Block1(nn.Module):
    def __init__(self, in_Channel, out_Channel):
        super(Basic_Block1, self).__init__()
        self.conv_x = nn.Sequential(
            nn.Conv2d(in_Channel, out_Channel, kernel_size=1, stride=2),
            nn.BatchNorm2d(out_Channel)
        )
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_Channel, out_Channel, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(out_Channel),
            nn.ReLU(True)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(out_Channel, out_Channel, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(out_Channel)
        )

    def forward(self, x):
        z = self.conv_x(x)
        y = F.relu(self.conv1(x))
        y = self.conv2(y)
        return F.relu(z+y)
